from bert_serving.client import BertClient
from sklearn.metrics.pairwise import cosine_similarity
from tool.bert_pooling_client import start_word_service
import threading
SIM_THREHOLD = 0.8
NEARLY_SIM_THREHOLD = 0.6


def read_data():
    """
    从词表中读取,返回包含所有单词的list
    :return:
    """
    with open("../data/words_clean.txt") as file:
        word_list = [line.strip() for line in file]
    # word_list = ["high", "low", "high", "legal"]
    return word_list


def compare_word(word_list):
    """
    1.调用consine 比较词与词之间的相似度,得到每个词与其他词的相似度矩阵
    2.过滤矩阵中 > 0.8的位置,存入map 中
    3.同时保存 > 0.7的位置,存入map 中
    4.将map 按行存储到文件中
    :return:
    """
    bc = BertClient()
    word_vecs = bc.encode(word_list)
    word_vecs = [vecs[1] for vecs in word_vecs]
    with open("../data/word_sim", 'w') as file:
        sim_vec = cosine_similarity(word_vecs, word_vecs)
    return sim_vec


def output(sim_vec,word_list):
    """
    根据sim_vec 过滤并输出到文件中
    :param sim_vec:
    :return:
    """
    sim_map = {}
    nearly_sim_map = {}
    for i in range(len(sim_vec)):
        sim_map[word_list[i]] = []
        nearly_sim_map[word_list[i]] = []
        for j in range(len(sim_vec[0])):
            if i == j:
                continue
            if sim_vec[i][j] > SIM_THREHOLD:
                sim_map[word_list[i]].append(word_list[j])
            elif sim_vec[i][j] > NEARLY_SIM_THREHOLD:
                nearly_sim_map[word_list[i]].append(word_list[j])
    with open("../data/sim_map", 'w') as sim_map_file:
        sim_map_file.writelines([key+" "+",".join(sim_map[key])+"\n" if len(sim_map[key]) > 0 else "" for key in sim_map])
    with open("../data/nearly_sim_map", 'w') as nearly_sim_map_file:
        nearly_sim_map_file.writelines([key+" "+",".join(nearly_sim_map[key])+"\n" if len(nearly_sim_map[key]) > 0 else "" for key in nearly_sim_map])


def main():
    word_list = read_data()
    sim_vec = compare_word(word_list)
    output(sim_vec, word_list)


if __name__ == "__main__":
    t = threading.Thread(target=start_word_service)
    t.start()
    main()
    print("cal finish")
    exit()
